#!/usr/bin/env python
# coding: utf8

"""
    Module that provides a class wrapper for source separation.

    :Example:

    >>> from segment.separator import Segmentor
    >>> separator = Segmentor('segment:2stems')
    >>> separator.segment(waveform, lambda instrument, data: ...)
    >>> separator.segment_to_file(...)
"""

import logging
import os
from multiprocessing import Pool
from os.path import basename, join, splitext

import numpy as np
import tensorflow as tf
from PIL import Image

from . import SegmentError
from .model import InputProviderFactory, EstimatorSpecBuilder
from .utils.configuration import load_configuration
from .utils.display import create_mask
from .utils.display import display
from .utils.estimator import create_estimator, to_predictor, get_default_model_dir
from .utils.logging import get_logger

logger = logging.getLogger("segment")


class Segmentor(object):
    """ A wrapper class for performing separation. """

    def __init__(self, params_descriptor, multiprocess=True):
        """ Default constructor.

        :param params_descriptor: Descriptor for TF params to be used.
        """
        self._params = load_configuration(params_descriptor)
        self._predictor = None
        self._input_provider = None
        self._builder = None
        self._features = None
        self._graph = None
        self._pool = Pool() if multiprocess else None
        self._tasks = []

    def _get_predictor(self):
        """ Lazy loading access method for internal predictor instance.

        :returns: Predictor to use for source separation.
        """
        if self._predictor is None:
            estimator = create_estimator(self._params)
            self._predictor = to_predictor(estimator)
        return self._predictor

    def join(self, timeout=200):
        """ Wait for all pending tasks to be finished.

        :param timeout: (Optional) task waiting timeout.
        """
        while len(self._tasks) > 0:
            task = self._tasks.pop()
            task.get()
            task.wait(timeout=timeout)

    def _segment_tensorflow(self, img_raw):
        """
        Performs source segment over the given image with tensorflow backend.

        :param img_raw: image to apply segment on.
        :returns: Separated waveforms.
        """
        predictor = self._get_predictor()
        prediction = predictor({'image_raw': img_raw})
        return prediction

    def _segment_session(self, img_raw):
        out = {}
        features = self._get_features()
        latest_checkpoint = tf.train.latest_checkpoint(get_default_model_dir(self._params['model_dir']))
        outputs = self._get_builder().outputs
        saver = tf.compat.v1.train.Saver()
        with tf.compat.v1.Session() as sess:
            saver.restore(sess, latest_checkpoint)
            feed_dict = self._get_input_provider().get_feed_dict(features, img_raw)
            outputs = sess.run(outputs, feed_dict=feed_dict)
            out['output'] = outputs
        return out

    def _segment_pb(self, img_raw):
        out = {}
        pb_file = join(get_default_model_dir(self._params['model_dir']), 'saved_model.pb')
        graph = self._get_graph(pb_file)
        input_node = graph.get_tensor_by_name("image_raw:0")
        output_node = graph.get_tensor_by_name("conv2d_18/truediv:0")
        with tf.compat.v1.Session(graph=graph) as sess:
            outputs = sess.run(output_node, feed_dict={input_node: img_raw})
            out['output'] = outputs
        return out

    def _get_input_provider(self):
        if self._input_provider is None:
            self._input_provider = InputProviderFactory.get(self._params)
        return self._input_provider

    def _get_features(self):
        if self._features is None:
            self._features = self._get_input_provider().get_input_dict_placeholders()
        return self._features

    def _get_builder(self):
        if self._builder is None:
            self._builder = EstimatorSpecBuilder(self._get_features(), self._params)
        return self._builder

    def _get_graph(self, pb_file):
        if self._graph is None:
            with open(pb_file, "rb") as fin:
                graph_def = tf.compat.v1.GraphDef()
                graph_def.ParseFromString(fin.read())
            with tf.compat.v1.Graph().as_default() as graph:
                tf.import_graph_def(
                    graph_def,
                    input_map=None,
                    return_elements=None,
                    name="",
                    producer_op_list=None
                )
            self._graph = graph
        return self._graph

    def _save(self, path, data):
        directory = os.path.dirname(path)
        if not os.path.exists(directory):
            raise SegmentError(f'output directory does not exists: {directory}')
        get_logger().debug('Writing file %s', path)
        img_data = np.asarray(data[0, :, :, 0] * 255, dtype=np.uint8)
        image = Image.fromarray(img_data)
        image.save(path)
        get_logger().info('File %s written succesfully', path)

    def segment(self, image_raw, image_descriptor=""):
        """ Performs separation on a waveform.

        :param image_raw:               image to be segment (as a numpy array)
        :param image_descriptor:    (Optional) string describing the waveform (e.g. filename).
        """
        if self._params["backend"] == "tensorflow":
            return self._segment_tensorflow(image_raw)
        elif self._params["backend"] == "pb":
            return self._segment_pb(image_raw)
        else:
            return self._segment_session(image_raw)

    def segment_to_file(self, image_descriptor, destination, synchronous=True):
        """ Performs source separation and export result to file using
        given audio adapter.

        Filename format should be a Python formattable string that could use
        following parameters : {instrument}, {filename} and {codec}.

        :param image_descriptor:    Describe image to segment, used by audio
                                    adapter to retrieve and load audio data,
                                    in case of file based audio adapter, such
                                    descriptor would be a file path.
        :param destination:         Target directory to write output to.
        :param synchronous:         (Optional) True is should by synchronous.
        """
        image = Image.open(image_descriptor)
        image = image.resize((self._params['H'], self._params['W']))
        img_raw = np.array(image)
        img_raw = img_raw / 255.0
        if not img_raw.dtype == b'float32':
            img_raw = img_raw.astype(b'float32')
        if img_raw.ndim == 2:
            img_raw = np.expand_dims(img_raw, -1)
        img_raw = np.expand_dims(img_raw, 0)
        sources = self.segment(img_raw, image_descriptor)
        self.save_to_file(sources, image_descriptor, destination, synchronous)
        mask = sources['output']
        if mask.ndim == 4 and mask.shape[-1] == 3:
            mask = create_mask(mask[0, :, :, :])
        else:
            mask = mask[0, :, :, :]
        display([img_raw[0, :, :, :], mask], (self._params['H'] / 100, self._params['W'] / 100))

    def save_to_file(
            self, sources, image_descriptor,
            destination,
            synchronous=True):
        """ export dictionary of sources to files.

        :param sources:             Dictionary of sources to be exported. The
                                    keys are the name of the instruments, and
                                    the values are Nx2 numpy arrays containing
                                    the corresponding intrument waveform, as
                                    returned by the separate method
        :param image_descriptor:    Describe song to separate, used by audio
                                    adapter to retrieve and load audio data,
                                    in case of file based audio adapter, such
                                    descriptor would be a file path.
        :param destination:         Target directory to write output to.
        :param synchronous:         (Optional) True is should by synchronous.

        """

        filename = splitext(basename(image_descriptor))[0]
        generated = []
        for instrument, data in sources.items():
            path = join(destination, '{filename}/{instrument}.jpg'.format(
                filename=filename,
                instrument=instrument))
            directory = os.path.dirname(path)
            if not os.path.exists(directory):
                os.makedirs(directory)
            if path in generated:
                raise SegmentError((
                    f'Separated source path conflict : {path},'
                    'please check your filename format'))
            generated.append(path)
            if self._pool:
                task = self._pool.apply_async(self._save, (path, data))
                self._tasks.append(task)
            else:
                self._save(path, data)
        if synchronous and self._pool:
            self.join()
